#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
@author: liang kang
@contact: gangkanli1219@gmail.com
@time: 2018/3/17 11:37
@desc: 
"""
import tensorflow as tf
from dltools.image.imaging_tf import histogram_equalization_tf
from dltools.data.input.models_input import DataAugmentDecoder
from object_detection.core import standard_fields as fields
from object_detection.protos import input_reader_pb2

parallel_reader = tf.contrib.slim.parallel_reader
slim_example_decoder = tf.contrib.slim.tfexample_decoder


class PeopleDataAugmentDecoder(DataAugmentDecoder):

    def __init__(self,
                 load_instance_masks=False,
                 instance_mask_type=input_reader_pb2.NUMERICAL_MASKS,
                 label_map_proto_file=None,
                 use_display_name=False,
                 dct_method=''):
        super(PeopleDataAugmentDecoder, self).__init__(
            load_instance_masks=load_instance_masks,
            instance_mask_type=instance_mask_type,
            label_map_proto_file=label_map_proto_file,
            use_display_name=use_display_name,
            dct_method=dct_method)

    def _data_augment(self, tensor_dict):
        image = tensor_dict[fields.InputDataFields.image]
        with tf.name_scope('HistogramEqualization', values=[image]):
            do_random = tf.greater(tf.random_uniform([]), 0.5)
            image = tf.cond(do_random, lambda: histogram_equalization_tf(image),
                            lambda: image)
        tensor_dict[fields.InputDataFields.image] = image


class CarDataAugmentDecoder(DataAugmentDecoder):

    def __init__(self,
                 load_instance_masks=False,
                 instance_mask_type=input_reader_pb2.NUMERICAL_MASKS,
                 label_map_proto_file=None,
                 use_display_name=False,
                 dct_method=''):
        super(CarDataAugmentDecoder, self).__init__(
            load_instance_masks=load_instance_masks,
            instance_mask_type=instance_mask_type,
            label_map_proto_file=label_map_proto_file,
            use_display_name=use_display_name,
            dct_method=dct_method)

    def _data_augment(self, tensor_dict):
        image = tensor_dict[fields.InputDataFields.image]
        with tf.name_scope('HistogramEqualization', values=[image]):
            do_random = tf.greater(tf.random_uniform([]), 0.5)
            image = tf.cond(do_random, lambda: histogram_equalization_tf(image),
                            lambda: image)
        tensor_dict[fields.InputDataFields.image] = image
